package first.group.ssh.hostkey;

import first.group.ssh.packet.reader.BigIntegerReader;
import first.group.ssh.packet.reader.StringReader;
import first.group.ssh.packet.writer.BigIntegerWriter;
import first.group.ssh.packet.writer.StringWriter;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.security.KeyFactory;
import java.security.PublicKey;
import java.security.Signature;
import java.security.spec.RSAPublicKeySpec;
import java.util.Arrays;
import java.util.logging.Logger;

public final class SshRsa extends HostKeyAlgorithm {

    @Override
    public String getName() {
        return "ssh-rsa";
    }

    public static class SshRsaKey extends HostKey<SshRsa> {
        public final BigInteger n, e;

        public SshRsaKey(BigInteger e, BigInteger n) {
            this.e = e;
            this.n = n;
        }

        @Override
        public SshRsa getAlgorithm() {
            return new SshRsa();
        }

        @Override
        public boolean check(byte[] data, HostSignature<? extends HostKeyAlgorithm> signature) {
            if (!(signature instanceof SshRsaSignature rsaSignature)) {
                throw new IllegalArgumentException("Invalid signature type.");
            }

            try {
                var sig = Signature.getInstance("SHA1withRSA");
                var spec = new RSAPublicKeySpec(this.n, this.e);
                var factory = KeyFactory.getInstance("RSA");
                sig.initVerify(factory.generatePublic(spec));
                sig.update(data);
                return sig.verify(rsaSignature.sign);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }

        @Override
        public void write(ByteBuffer buffer) {
            var posStart = buffer.position();
            buffer.position(buffer.position() + 4);
            
            var dataStart = buffer.position();
            StringWriter.writeString(buffer, getAlgorithm().getName());
            BigIntegerWriter.writeBigInteger(buffer, e);
            BigIntegerWriter.writeBigInteger(buffer, n);
            
            var dataEnd = buffer.position();
            buffer.position(posStart);
            buffer.putInt(dataEnd - dataStart);
            buffer.position(dataEnd);
        }

        @Override
        public String toString() {
            return "SshRsaKey{" + "n=" + n + ", e=" + e + '}';
        }
    }

    private static class SshRsaKeyParser extends HostKeyParser<SshRsa> {

        @Override
        public HostKey<SshRsa> parse(ByteBuffer buffer) {
            var e = BigIntegerReader.readBigInteger(buffer);
            if (e == null) {
                return null;
            }
            var n = BigIntegerReader.readBigInteger(buffer);
            if (n == null) {
                return null;
            }
            return new SshRsaKey(e, n);
        }
    }
    
    private static final SshRsaKeyParser keyParser = new SshRsaKeyParser();

    @Override
    public HostKeyParser<SshRsa> getKeyParser() {
        return keyParser;
    }

    public static class SshRsaSignature extends HostSignature<SshRsa> {
        public final byte[] sign;

        public SshRsaSignature(byte[] sign) {
            this.sign = sign;
        }

        @Override
        public SshRsa getAlgorithm() {
            return new SshRsa();
        }

        @Override
        public String toString() {
            return "SshRsaSignature{" + "sign=" + Arrays.toString(sign) + '}';
        }
    }

    private static class SshRsaSignatureParser extends HostSignatureParser<SshRsa> {

        @Override
        public HostSignature<SshRsa> parse(ByteBuffer buffer) {
            var sign = StringReader.readByteArray(buffer);
            if (sign == null) {
                return null;
            }
            return new SshRsaSignature(sign);
        }
    }
    
    private static final SshRsaSignatureParser signatureParser = new SshRsaSignatureParser();

    @Override
    public HostSignatureParser<SshRsa> getSignatureParser() {
        return signatureParser;
    }
}
